# Copyright 2025 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""expected output of MLP"""
import numpy as np
import mindspore as ms
from mindformers.parallel_core.transformer_config import TransformerConfig


def get_init_params(args, config: TransformerConfig, seq_length=4, batch_size=2):
    """Generate initialization parameters"""
    np.random.seed(1)
    input_size = args.input_size if args.input_size else config.hidden_size
    if config.gated_linear_unit:
        state_dict = {
            "mlp.linear_fc1.weight": 0.01 * np.random.rand(2 * config.ffn_hidden_size, input_size),
            "mlp.linear_fc2.weight": 0.01 * np.random.rand(config.hidden_size, config.ffn_hidden_size)
        }
        if config.tensor_model_parallel_size == 2:
            part_0, part_1, part_2, part_3 = np.split(state_dict['mlp.linear_fc1.weight'], 4, axis=0)
            state_dict['mlp.linear_fc1.weight'] = np.concatenate([part_0, part_2, part_1, part_3], axis=0)
    else:
        state_dict = {
            "mlp.linear_fc1.weight": 0.01 * np.random.rand(config.ffn_hidden_size, input_size),
            "mlp.linear_fc2.weight": 0.01 * np.random.rand(config.hidden_size, config.ffn_hidden_size)
        }
    for k in state_dict:
        state_dict[k] = ms.Parameter(ms.tensor(state_dict[k], dtype=ms.float32))

    input_ = np.random.rand(seq_length, batch_size, input_size)
    input_ = ms.tensor(input_, dtype=ms.bfloat16)

    return input_, state_dict


def get_golden_datas(args) -> np.ndarray:
    """Generate golden data for test."""
    if args.input_size == 32 is not None:
        output_golden = np.array(
            [[[0.0002843735, 0.0003589125, 0.0003017838, 0.0003067734,
               0.0003109106, 0.0002981229, 0.0002762848, 0.0003798283,
               0.0002844124, 0.0002834244, 0.0003677915, 0.0003187459,
               0.0003219039, 0.0002767901, 0.0003016057, 0.0002498944],
              [0.0002506988, 0.0003244758, 0.0002593943, 0.0002691273,
               0.0002715064, 0.0002678589, 0.0002447706, 0.0003291702,
               0.0002432512, 0.0002446361, 0.0003225597, 0.0002730018,
               0.0002809532, 0.0002455753, 0.0002583148, 0.0002182310]],
             [[0.0003119569, 0.0004044899, 0.0003295057, 0.0003350345,
               0.0003449772, 0.0003258427, 0.0003029937, 0.0004144278,
               0.0003144512, 0.0003187756, 0.0004083854, 0.0003471685,
               0.0003455959, 0.0003007032, 0.0003314289, 0.0002745171],
              [0.0003268698, 0.0004171443, 0.0003380423, 0.0003473255,
               0.0003556831, 0.0003446467, 0.0003163249, 0.0004299119,
               0.0003194707, 0.0003200345, 0.0004169832, 0.0003525523,
               0.0003518050, 0.0003163963, 0.0003438482, 0.0002806541]],
             [[0.0002283673, 0.0002945957, 0.0002480418, 0.0002498238,
               0.0002528055, 0.0002393862, 0.0002205361, 0.0003087847,
               0.0002295643, 0.0002287964, 0.0003007228, 0.0002591236,
               0.0002584554, 0.0002202614, 0.0002478469, 0.0002010728],
              [0.0001893662, 0.0002481799, 0.0001996867, 0.0002013757,
               0.0002069773, 0.0001973989, 0.0001852193, 0.0002522317,
               0.0001864916, 0.0001951331, 0.0002464133, 0.0002084459,
               0.0002095294, 0.0001841778, 0.0001982806, 0.0001639207]],
             [[0.0002806637, 0.0003573631, 0.0002937675, 0.0003028220,
               0.0003065119, 0.0003010270, 0.0002730124, 0.0003712290,
               0.0002734405, 0.0002786421, 0.0003623499, 0.0003114895,
               0.0003216742, 0.0002781576, 0.0002872665, 0.0002481540],
              [0.0002061773, 0.0002699204, 0.0002196116, 0.0002186791,
               0.0002230823, 0.0002154907, 0.0002004602, 0.0002743427,
               0.0002081957, 0.0002078231, 0.0002681023, 0.0002244467,
               0.0002327584, 0.0002099627, 0.0002148721, 0.0001803946]]], dtype=np.float32)
    elif args.gated_linear_unit:
        output_golden = np.array(
            [[[0.0000733683, 0.0000909508, 0.0000698835, 0.0000924880,
               0.0000818640, 0.0000622055, 0.0000739467, 0.0000799153,
               0.0000613963, 0.0000783803, 0.0001000925, 0.0000622146,
               0.0000825421, 0.0000920034, 0.0000837915, 0.0000794538],
              [0.0000712826, 0.0000899131, 0.0000682673, 0.0000903676,
               0.0000802998, 0.0000608406, 0.0000726998, 0.0000809333,
               0.0000613315, 0.0000799796, 0.0000957279, 0.0000623976,
               0.0000820426, 0.0000919889, 0.0000813653, 0.0000806986]],
             [[0.0000715935, 0.0000869136, 0.0000652340, 0.0000865586,
               0.0000756654, 0.0000545406, 0.0000684799, 0.0000776089,
               0.0000568416, 0.0000765296, 0.0000909929, 0.0000576520,
               0.0000780936, 0.0000876632, 0.0000774766, 0.0000762414],
              [0.0000565747, 0.0000683313, 0.0000524830, 0.0000693377,
               0.0000612595, 0.0000447933, 0.0000551950, 0.0000600794,
               0.0000456970, 0.0000579404, 0.0000745601, 0.0000464607,
               0.0000611514, 0.0000688959, 0.0000618574, 0.0000586497]],
             [[0.0000475586, 0.0000567166, 0.0000434163, 0.0000556350,
               0.0000505326, 0.0000368306, 0.0000435420, 0.0000520190,
               0.0000373207, 0.0000498961, 0.0000599300, 0.0000388123,
               0.0000510283, 0.0000581764, 0.0000494977, 0.0000496677],
              [0.0000790294, 0.0000975368, 0.0000734458, 0.0000979018,
               0.0000879520, 0.0000647741, 0.0000777439, 0.0000885753,
               0.0000676341, 0.0000855171, 0.0001044117, 0.0000675433,
               0.0000888914, 0.0001019846, 0.0000868552, 0.0000866440]],
             [[0.0000510057, 0.0000659857, 0.0000510755, 0.0000665491,
               0.0000574119, 0.0000449925, 0.0000532988, 0.0000589157,
               0.0000445915, 0.0000595711, 0.0000696541, 0.0000439622,
               0.0000604231, 0.0000656269, 0.0000605255, 0.0000585135],
              [0.0000542883, 0.0000657589, 0.0000502881, 0.0000662970,
               0.0000581312, 0.0000414027, 0.0000525229, 0.0000580756,
               0.0000438181, 0.0000575348, 0.0000698252, 0.0000446309,
               0.0000585120, 0.0000660210, 0.0000589243, 0.0000574732]]], dtype=np.float32)
    else:
        output_golden = np.array(
            [[[0.0017414739, 0.0018487792, 0.0017199146, 0.0013846543,
               0.0012862221, 0.0020816431, 0.0019175113, 0.0014373211,
               0.0013844393, 0.0015804252, 0.0012110893, 0.0018034461,
               0.0014611555, 0.0017085664, 0.0015952593, 0.0015352677],
              [0.0019919416, 0.0020969177, 0.0019119657, 0.0016831061,
               0.0015097546, 0.0024931964, 0.0022334943, 0.0015607530,
               0.0016224522, 0.0018226021, 0.0013644423, 0.0020328050,
               0.0016648017, 0.0019004948, 0.0018370659, 0.0017043886]],
             [[0.0015751260, 0.0016880741, 0.0015236203, 0.0013121577,
               0.0011900065, 0.0019737678, 0.0017540415, 0.0012518927,
               0.0012730283, 0.0014253005, 0.0010885255, 0.0016146263,
               0.0013117085, 0.0015241898, 0.0014786809, 0.0013727419],
              [0.0019514252, 0.0020915624, 0.0018961922, 0.0016189103,
               0.0014609485, 0.0024405473, 0.0021871955, 0.0015791889,
               0.0015751384, 0.0017856713, 0.0013288374, 0.0019960622,
               0.0016252893, 0.0018935392, 0.0017961639, 0.0017028205]],
             [[0.0018867148, 0.0020168365, 0.0018601861, 0.0015550752,
               0.0014348353, 0.0022915448, 0.0021151179, 0.0015194945,
               0.0015332238, 0.0017510453, 0.0013417912, 0.0019695892,
               0.0015768812, 0.0018295512, 0.0017305955, 0.0016419239],
              [0.0014538650, 0.0015827805, 0.0014461200, 0.0012378941,
               0.0010918338, 0.0018326158, 0.0016430839, 0.0011372042,
               0.0012296749, 0.0013730322, 0.0010569921, 0.0015125085,
               0.0011989862, 0.0014621870, 0.0013702807, 0.0012903372]],
             [[0.0017733796, 0.0019430587, 0.0017720656, 0.0014609765,
               0.0013720200, 0.0021825042, 0.0020178643, 0.0014793403,
               0.0014748116, 0.0016822204, 0.0012621277, 0.0018532595,
               0.0014871625, 0.0017542860, 0.0016103230, 0.0015683112],
              [0.0018832930, 0.0020078733, 0.0018468696, 0.0015462689,
               0.0014222845, 0.0022924296, 0.0021196704, 0.0015214478,
               0.0015399234, 0.0017200317, 0.0013087805, 0.0019487045,
               0.0015966668, 0.0018407278, 0.0016958695, 0.0016332783]]], dtype=np.float32)

    return output_golden


def get_gpu_datas(args) -> np.ndarray:
    """Generate gpu data for test."""
    if args.input_size == 32:
        output_gpu = np.array(
            [[[0.0002841949, 0.0003585815, 0.0003013611, 0.0003070831,
               0.0003108978, 0.0002975464, 0.0002765656, 0.0003795624,
               0.0002841949, 0.0002841949, 0.0003681183, 0.0003204346,
               0.0003223419, 0.0002765656, 0.0003013611, 0.0002498627],
              [0.0002517700, 0.0003242493, 0.0002593994, 0.0002689362,
               0.0002708435, 0.0002670288, 0.0002441406, 0.0003299713,
               0.0002431870, 0.0002441406, 0.0003223419, 0.0002727509,
               0.0002803802, 0.0002460480, 0.0002593994, 0.0002183914]],
             [[0.0003128052, 0.0004043579, 0.0003299713, 0.0003356934,
               0.0003452301, 0.0003261566, 0.0003032684, 0.0004138947,
               0.0003147125, 0.0003185272, 0.0004081726, 0.0003490448,
               0.0003471375, 0.0003013611, 0.0003318787, 0.0002746582],
              [0.0003261566, 0.0004158020, 0.0003376007, 0.0003471375,
               0.0003547668, 0.0003452301, 0.0003166199, 0.0004291534,
               0.0003185272, 0.0003204346, 0.0004177094, 0.0003528595,
               0.0003528595, 0.0003166199, 0.0003433228, 0.0002803802]],
             [[0.0002288818, 0.0002956390, 0.0002479553, 0.0002498627,
               0.0002536774, 0.0002393723, 0.0002202988, 0.0003089905,
               0.0002298355, 0.0002288818, 0.0003013611, 0.0002593994,
               0.0002593994, 0.0002212524, 0.0002479553, 0.0002012253],
              [0.0001888275, 0.0002479553, 0.0002002716, 0.0002012253,
               0.0002069473, 0.0001974106, 0.0001850128, 0.0002517700,
               0.0001859665, 0.0001955032, 0.0002460480, 0.0002088547,
               0.0002098083, 0.0001840591, 0.0001983643, 0.0001640320]],
             [[0.0002803802, 0.0003566742, 0.0002937317, 0.0003032684,
               0.0003070831, 0.0003013611, 0.0002727509, 0.0003719330,
               0.0002727509, 0.0002784729, 0.0003623962, 0.0003128052,
               0.0003223419, 0.0002784729, 0.0002880096, 0.0002479553],
              [0.0002059937, 0.0002689362, 0.0002193451, 0.0002183914,
               0.0002231598, 0.0002155304, 0.0002002716, 0.0002746582,
               0.0002079010, 0.0002079010, 0.0002689362, 0.0002241135,
               0.0002326965, 0.0002107620, 0.0002145767, 0.0001802444]]], dtype=np.float16)
    elif args.gated_linear_unit:
        output_gpu = np.array(
            [[[0.0000734329, 0.0000910759, 0.0000700951, 0.0000925064,
               0.0000820160, 0.0000624657, 0.0000739098, 0.0000801086,
               0.0000615120, 0.0000786781, 0.0001001358, 0.0000624657,
               0.0000824928, 0.0000920296, 0.0000839233, 0.0000796318],
              [0.0000715256, 0.0000901222, 0.0000681877, 0.0000905991,
               0.0000805855, 0.0000610352, 0.0000729561, 0.0000810623,
               0.0000615120, 0.0000801086, 0.0000958443, 0.0000624657,
               0.0000824928, 0.0000920296, 0.0000815392, 0.0000810623]],
             [[0.0000715256, 0.0000867844, 0.0000653267, 0.0000867844,
               0.0000758171, 0.0000545979, 0.0000686646, 0.0000777245,
               0.0000567436, 0.0000767708, 0.0000910759, 0.0000576973,
               0.0000782013, 0.0000877380, 0.0000772476, 0.0000762939],
              [0.0000562668, 0.0000681877, 0.0000522137, 0.0000691414,
               0.0000610352, 0.0000445843, 0.0000550747, 0.0000598431,
               0.0000455379, 0.0000576973, 0.0000743866, 0.0000462532,
               0.0000610352, 0.0000686646, 0.0000615120, 0.0000584126]],
             [[0.0000474453, 0.0000567436, 0.0000433922, 0.0000555515,
               0.0000505447, 0.0000367165, 0.0000436306, 0.0000519753,
               0.0000371933, 0.0000498295, 0.0000598431, 0.0000388622,
               0.0000510216, 0.0000581741, 0.0000495911, 0.0000495911],
              [0.0000791550, 0.0000972748, 0.0000734329, 0.0000977516,
               0.0000877380, 0.0000648499, 0.0000777245, 0.0000886917,
               0.0000677109, 0.0000853539, 0.0001044273, 0.0000677109,
               0.0000891685, 0.0001020432, 0.0000867844, 0.0000867844]],
             [[0.0000510216, 0.0000658035, 0.0000510216, 0.0000667572,
               0.0000574589, 0.0000450611, 0.0000531673, 0.0000588894,
               0.0000445843, 0.0000596046, 0.0000696182, 0.0000438690,
               0.0000605583, 0.0000658035, 0.0000605583, 0.0000584126],
              [0.0000543594, 0.0000658035, 0.0000503063, 0.0000662804,
               0.0000581741, 0.0000414848, 0.0000526905, 0.0000581741,
               0.0000438690, 0.0000576973, 0.0000700951, 0.0000445843,
               0.0000586510, 0.0000662804, 0.0000591278, 0.0000574589]]], dtype=np.float16)
    else:
        output_gpu = np.array(
            [[[0.0017395020, 0.0018463135, 0.0017166138, 0.0013809204,
               0.0012817383, 0.0020751953, 0.0019149780, 0.0014343262,
               0.0013809204, 0.0015792847, 0.0012054443, 0.0018005371,
               0.0014572144, 0.0017089844, 0.0015945435, 0.0015335083],
              [0.0019989014, 0.0020904541, 0.0019149780, 0.0016860962,
               0.0015106201, 0.0024871826, 0.0022277832, 0.0015563965,
               0.0016174316, 0.0018234253, 0.0013656616, 0.0020294189,
               0.0016632080, 0.0018997192, 0.0018386841, 0.0017013550]],
             [[0.0015716553, 0.0016860962, 0.0015258789, 0.0013122559,
               0.0011901855, 0.0019683838, 0.0017547607, 0.0012512207,
               0.0012741089, 0.0014266968, 0.0010910034, 0.0016174316,
               0.0013122559, 0.0015258789, 0.0014801025, 0.0013732910],
              [0.0019531250, 0.0020904541, 0.0018997192, 0.0016174316,
               0.0014572144, 0.0024414062, 0.0021820068, 0.0015792847,
               0.0015792847, 0.0017852783, 0.0013275146, 0.0019989014,
               0.0016250610, 0.0018920898, 0.0017929077, 0.0017013550]],
             [[0.0018844604, 0.0020141602, 0.0018615723, 0.0015563965,
               0.0014343262, 0.0022888184, 0.0021209717, 0.0015182495,
               0.0015335083, 0.0017471313, 0.0013427734, 0.0019683838,
               0.0015716553, 0.0018310547, 0.0017318726, 0.0016403198],
              [0.0014572144, 0.0015869141, 0.0014495850, 0.0012359619,
               0.0010910034, 0.0018310547, 0.0016479492, 0.0011367798,
               0.0012283325, 0.0013732910, 0.0010604858, 0.0015106201,
               0.0011978149, 0.0014648438, 0.0013732910, 0.0012893677]],
             [[0.0017776489, 0.0019454956, 0.0017700195, 0.0014648438,
               0.0013732910, 0.0021820068, 0.0020141602, 0.0014801025,
               0.0014724731, 0.0016784668, 0.0012588501, 0.0018539429,
               0.0014877319, 0.0017547607, 0.0016098022, 0.0015640259],
              [0.0018844604, 0.0020141602, 0.0018463135, 0.0015487671,
               0.0014190674, 0.0022888184, 0.0021209717, 0.0015182495,
               0.0015411377, 0.0017166138, 0.0013046265, 0.0019454956,
               0.0015945435, 0.0018386841, 0.0016937256, 0.0016326904]]], dtype=np.float16)

    return output_gpu


def get_grads(device) -> list[np.ndarray]:
    """Generate grads data for test."""
    if device == 'cpu':
        output = [np.array([[[0.00358849, 0.00397592, 0.00281392, 0.00408467,
                              0.00329911, 0.00313804, 0.00287723, 0.00327443,
                              0.00397678, 0.00393216, 0.00253552, 0.00414877,
                              0.00264414, 0.00396701, 0.00301837, 0.00362666],
                             [0.00361133, 0.00400056, 0.00282536, 0.00409804,
                              0.00332681, 0.00316927, 0.00289322, 0.00328997,
                              0.00399494, 0.00395481, 0.00253725, 0.00416637,
                              0.00266053, 0.00399461, 0.00303685, 0.00364819]],
                            [[0.00357625, 0.00396229, 0.00280036, 0.00406456,
                              0.0032961, 0.00313534, 0.00286683, 0.00326088,
                              0.00396103, 0.00391977, 0.00251818, 0.00413285,
                              0.00263667, 0.00395849, 0.00301202, 0.00361661],
                             [0.00360544, 0.00399649, 0.00282541, 0.00409751,
                              0.00332099, 0.00316221, 0.00288866, 0.00328947,
                              0.00399377, 0.00394908, 0.00253792, 0.00416786,
                              0.00265755, 0.00398979, 0.00303717, 0.0036475]],
                            [[0.00360087, 0.00399243, 0.00281884, 0.00409584,
                              0.00331547, 0.00315367, 0.00288709, 0.00328134,
                              0.00398511, 0.00394567, 0.00254016, 0.00415713,
                              0.00265323, 0.00398386, 0.00302817, 0.00364326],
                             [0.00357491, 0.00395741, 0.00279007, 0.00405445,
                              0.00328517, 0.00312859, 0.0028571, 0.00325411,
                              0.00394988, 0.00391398, 0.00251166, 0.00412257,
                              0.00262869, 0.00395389, 0.00300256, 0.00361064]],
                            [[0.003597, 0.00398737, 0.00281334, 0.00408857,
                              0.00330749, 0.003149, 0.00287716, 0.0032785,
                              0.00397591, 0.00393754, 0.00253427, 0.00414818,
                              0.00264593, 0.00397641, 0.00302268, 0.00363965],
                             [0.00360494, 0.00399137, 0.00282102, 0.00409681,
                              0.00331309, 0.00315501, 0.00288683, 0.00328423,
                              0.00398583, 0.00394463, 0.00254061, 0.00415578,
                              0.00265328, 0.00398023, 0.00302891, 0.00364053]]], dtype=np.float32),
                  np.array([[0.11603601, 0.12541088, 0.2167649, 0.13615835,
                             0.23204474, 0.14781404, 0.14284374, 0.21878365,
                             0.2296954, 0.14339207, 0.15646756, 0.14679496,
                             0.16158302, 0.19834986, 0.15242273, 0.14654386],
                            [0.11827907, 0.1275674, 0.22092944, 0.13856976,
                             0.23650697, 0.15069492, 0.14558686, 0.22275625,
                             0.23401098, 0.1460322, 0.15909217, 0.14930482,
                             0.16468957, 0.2020081, 0.15535495, 0.1491492],
                            [0.13304205, 0.14345269, 0.2486869, 0.15589577,
                             0.2661202, 0.1696076, 0.16383663, 0.25082958,
                             0.2634721, 0.16432543, 0.179159, 0.16816773,
                             0.18555485, 0.22716357, 0.1748758, 0.16771781],
                            [0.10260423, 0.11075466, 0.191801, 0.12039334,
                             0.20528972, 0.13070573, 0.12639251, 0.19352251,
                             0.20330656, 0.12684713, 0.13845786, 0.1298268,
                             0.14307016, 0.17528139, 0.13485247, 0.12945984],
                            [0.11434869, 0.12354562, 2141.9811*1e-4, 0.13414538,
                             0.22926822, 0.1459527, 0.14123122, 0.21591584,
                             0.2267648, 0.14131913, 0.15428637, 0.14478159,
                             0.15969388, 0.19552092, 0.15044016, 0.14439406],
                            [0.120708, 0.1303881, 0.2259141, 0.14160272,
                             0.24170484, 0.15409063, 0.14886515, 0.2280002,
                             0.23921087, 0.1491328, 0.1628059, 0.15289044,
                             0.16856587, 0.20621437, 0.158703, 0.15232202],
                            [0.10878122, 0.11732533, 0.20328, 0.12746143,
                             0.21753345, 0.13861288, 0.13398571, 0.20500873,
                             0.2153935, 0.1343977, 0.1465252, 0.13749526,
                             0.15160887, 0.18574253, 0.14290352, 0.13716751],
                            [0.11188044, 0.12077136, 0.20890741, 0.13109268,
                             0.2236695, 0.14262542, 0.13761999, 0.21071184,
                             0.22118692, 0.1379311, 0.15043199, 0.14113279,
                             0.15575643, 0.19107178, 0.14689873, 0.14116006],
                            [0.11597929, 0.1254121, 0.21746695, 0.13651237,
                             0.23271666, 0.14797877, 0.1434417, 0.21954595,
                             0.230406, 0.14357759, 0.15727311, 0.14724942,
                             0.16227552, 198.07822*1e-3, 0.15245219, 0.14640447],
                            [0.08209501, 0.08879905, 0.1537293, 0.09646383,
                             0.16456051, 0.10477749, 0.10125058, 0.15519002,
                             0.16279654, 0.10150903, 0.11086444, 0.10409537,
                             0.11466062, 0.14046507, 0.10806718, 0.10372541],
                            [0.12437321, 0.1344379, 0.2327614, 0.14611176,
                             0.24909762, 0.15863675, 0.15335996, 0.23501459,
                             0.24660684, 0.15379256, 0.16805252, 0.15765813,
                             0.17366807, 0.21256529, 0.16354786, 0.15703209],
                            [0.09975651, 0.10756143, 0.18641558, 0.1168772,
                             0.1995354, 0.12704614, 0.12287364, 0.18791163,
                             0.1975441, 0.12331721, 0.13431014, 0.12605324,
                             0.13895896, 0.17042944, 0.1311049, 0.12579435],
                            [0.1312534, 0.14187193, 0.24563143, 0.15428098,
                             0.26285297, 0.16739973, 0.16193543, 0.24805431,
                             0.26024646, 0.16220912, 0.1775631, 0.16630213,
                             0.18331951, 0.22403438, 0.17235264, 0.165645],
                            [0.13379085, 0.14438005, 0.24935594, 0.15713774,
                             0.2670285, 0.17011712, 0.16429949, 0.25189108,
                             0.2643592, 0.16508101, 0.18032789, 0.16869555,
                             0.18594846, 0.22805361, 0.17518404, 0.1686182],
                            [0.13005687, 0.14079395, 0.24347872, 0.15281783,
                             0.26067483, 0.16601822, 0.16050401, 0.24573143,
                             0.25768995, 0.16061778, 0.17561646, 0.1647053,
                             0.18148333, 0.22239792, 0.1709137, 0.16438828],
                            [0.12003458, 0.12975979, 0.22462347, 0.14066692,
                             0.24042274, 0.15318915, 0.14802292, 0.22646685,
                             0.23779865, 0.14833215, 0.16165313, 0.15197855,
                             0.16739687, 0.20548022, 0.15800232, 0.15169764]], dtype=np.float32),
                  np.array([[0.11688521, 0.16022952, 0.19401, 0.11149846,
                             0.17232549, 0.18658653, 0.17064336, 0.15132627,
                             0.14905359, 0.15008312, 0.15314576, 0.13133213,
                             0.20710704, 0.13642213, 0.19535986, 0.19016215],
                            [0.11688521, 0.16022952, 0.19401, 0.11149846,
                             0.17232549, 0.18658653, 0.17064336, 0.15132627,
                             0.14905359, 0.15008312, 0.15314576, 0.13133213,
                             0.20710704, 0.13642213, 0.19535986, 0.19016215],
                            [0.11688521, 0.16022952, 0.19401, 0.11149846,
                             0.17232549, 0.18658653, 0.17064336, 0.15132627,
                             0.14905359, 0.15008312, 0.15314576, 0.13133213,
                             0.20710704, 0.13642213, 0.19535986, 0.19016215],
                            [0.11688521, 0.16022952, 0.19401, 0.11149846,
                             0.17232549, 0.18658653, 0.17064336, 0.15132627,
                             0.14905359, 0.15008312, 0.15314576, 0.13133213,
                             0.20710704, 0.13642213, 0.19535986, 0.19016215],
                            [0.11688521, 0.16022952, 0.19401, 0.11149846,
                             0.17232549, 0.18658653, 0.17064336, 0.15132627,
                             0.14905359, 0.15008312, 0.15314576, 0.13133213,
                             0.20710704, 0.13642213, 0.19535986, 0.19016215],
                            [0.11688521, 0.16022952, 0.19401, 0.11149846,
                             0.17232549, 0.18658653, 0.17064336, 0.15132627,
                             0.14905359, 0.15008312, 0.15314576, 0.13133213,
                             0.20710704, 0.13642213, 0.19535986, 0.19016215],
                            [0.11688521, 0.16022952, 0.19401, 0.11149846,
                             0.17232549, 0.18658653, 0.17064336, 0.15132627,
                             0.14905359, 0.15008312, 0.15314576, 0.13133213,
                             0.20710704, 0.13642213, 0.19535986, 0.19016215],
                            [0.11688521, 0.16022952, 0.19401, 0.11149846,
                             0.17232549, 0.18658653, 0.17064336, 0.15132627,
                             0.14905359, 0.15008312, 0.15314576, 0.13133213,
                             0.20710704, 0.13642213, 0.19535986, 0.19016215],
                            [0.11688521, 0.16022952, 0.19401, 0.11149846,
                             0.17232549, 0.18658653, 0.17064336, 0.15132627,
                             0.14905359, 0.15008312, 0.15314576, 0.13133213,
                             0.20710704, 0.13642213, 0.19535986, 0.19016215],
                            [0.11688521, 0.16022952, 0.19401, 0.11149846,
                             0.17232549, 0.18658653, 0.17064336, 0.15132627,
                             0.14905359, 0.15008312, 0.15314576, 0.13133213,
                             0.20710704, 0.13642213, 0.19535986, 0.19016215],
                            [0.11688521, 0.16022952, 0.19401, 0.11149846,
                             0.17232549, 0.18658653, 0.17064336, 0.15132627,
                             0.14905359, 0.15008312, 0.15314576, 0.13133213,
                             0.20710704, 0.13642213, 0.19535986, 0.19016215],
                            [0.11688521, 0.16022952, 0.19401, 0.11149846,
                             0.17232549, 0.18658653, 0.17064336, 0.15132627,
                             0.14905359, 0.15008312, 0.15314576, 0.13133213,
                             0.20710704, 0.13642213, 0.19535986, 0.19016215],
                            [0.11688521, 0.16022952, 0.19401, 0.11149846,
                             0.17232549, 0.18658653, 0.17064336, 0.15132627,
                             0.14905359, 0.15008312, 0.15314576, 0.13133213,
                             0.20710704, 0.13642213, 0.19535986, 0.19016215],
                            [0.11688521, 0.16022952, 0.19401, 0.11149846,
                             0.17232549, 0.18658653, 0.17064336, 0.15132627,
                             0.14905359, 0.15008312, 0.15314576, 0.13133213,
                             0.20710704, 0.13642213, 0.19535986, 0.19016215],
                            [0.11688521, 0.16022952, 0.19401, 0.11149846,
                             0.17232549, 0.18658653, 0.17064336, 0.15132627,
                             0.14905359, 0.15008312, 0.15314576, 0.13133213,
                             0.20710704, 0.13642213, 0.19535986, 0.19016215],
                            [0.11688521, 0.16022952, 0.19401, 0.11149846,
                             0.17232549, 0.18658653, 0.17064336, 0.15132627,
                             0.14905359, 0.15008312, 0.15314576, 0.13133213,
                             0.20710704, 0.13642213, 0.19535986, 0.19016215]],
                           dtype=np.float32)]
    elif device == 'gpu':
        output = [np.array([[[0.003586, 0.003967, 0.002808, 0.00409, 0.003296, 0.003143, 0.002869, 0.003265,
                              0.003967, 0.003937, 0.002533, 0.00415, 0.00264, 0.003967, 0.003021, 0.003632],
                             [0.003616, 0.003998, 0.002823, 0.00409, 0.003326, 0.003174, 0.0029, 0.00328,
                              0.003998, 0.003967, 0.002533, 0.00415, 0.002655, 0.003998, 0.003036, 0.003647]],
                            [[0.003586, 0.003967, 0.002808, 0.00406, 0.003296, 0.003143, 0.002869, 0.003265,
                              0.003967, 0.003906, 0.002518, 0.00412, 0.00264, 0.003967, 0.003006, 0.003616],
                             [0.003616, 0.003998, 0.002823, 0.00409, 0.003326, 0.003159, 0.002884, 0.00328,
                              0.003998, 0.003937, 0.002533, 0.00415, 0.002655, 0.003998, 0.003036, 0.003647]],
                            [[0.003601, 0.003998, 0.002823, 0.00409, 0.003311, 0.003159, 0.002884, 0.00328,
                              0.003998, 0.003937, 0.002548, 0.00415, 0.002655, 0.003998, 0.003036, 0.003647],
                             [0.00357, 0.003967, 0.002792, 0.00406, 0.00328, 0.003128, 0.002853, 0.00325,
                              0.003937, 0.003906, 0.002518, 0.00412, 0.002625, 0.003967, 0.003006, 0.003616]],
                            [[0.003601, 0.003998, 0.002808, 0.00409, 0.003311, 0.003159, 0.002884, 0.00328,
                              0.003967, 0.003937, 0.002533, 0.00415, 0.00264, 0.003967, 0.003021, 0.003647],
                             [0.003616, 0.003998, 0.002823, 0.00409, 0.003311, 0.003159, 0.002884, 0.00328,
                              0.003998, 0.003937, 0.002533, 0.00415, 0.002655, 0.003998, 0.003036, 0.003647]]],
                           dtype=np.float16),
                  np.array([[0.1157, 0.125, 0.2168, 0.1357, 0.2314, 0.1475, 0.1426, 0.2188,
                             0.2295, 0.1436, 0.1562, 0.1465, 0.1611, 0.1982, 0.1523, 0.1465],
                            [0.11865, 0.1279, 0.2217, 0.1387, 0.2373, 0.1514, 0.1455, 0.2236,
                             0.2344, 0.1465, 0.1592, 0.1494, 0.165, 0.2021, 0.1553, 0.1494],
                            [0.1328, 0.1436, 0.248, 0.1553, 0.2656, 0.169, 0.1631, 0.25,
                             0.2637, 0.1641, 0.1787, 0.168, 0.1855, 0.2266, 0.1738, 0.167],
                            [0.10254, 0.11084, 0.1914, 0.1206, 0.2051, 0.1309, 0.126, 0.1934,
                             0.2031, 0.127, 0.1387, 0.1299, 0.1426, 0.1758, 0.1348, 0.1299],
                            [0.11426, 0.12354, 0.2148, 0.1338, 0.2295, 0.1465, 0.1416, 0.2158,
                             0.2266, 0.1416, 0.1543, 0.1445, 0.1602, 0.1963, 0.1504, 0.1445],
                            [0.1201, 0.1299, 0.2256, 0.1416, 0.2412, 0.1533, 0.1484, 0.2275,
                             0.2383, 0.1484, 0.1621, 0.1523, 0.168, 0.206, 0.1582, 0.1523],
                            [0.1089, 0.1177, 0.2041, 0.1279, 0.2178, 0.1387, 0.1338, 0.2051,
                             0.2158, 0.1348, 0.1465, 0.1377, 0.1523, 0.1865, 0.1436, 0.1377],
                            [0.1123, 0.1211, 0.209, 0.1318, 0.2236, 0.1426, 0.1377, 0.2109,
                             0.2217, 0.1377, 0.1504, 0.1416, 0.1562, 0.1914, 0.1475, 0.1416],
                            [0.1157, 0.125, 0.2168, 0.1367, 0.2324, 0.1475, 0.1436, 0.2188,
                             0.2305, 0.1436, 0.1572, 0.1475, 0.1621, 0.1982, 0.1523, 0.1465],
                            [0.08203, 0.0884, 0.1533, 0.0962, 0.1641, 0.1045, 0.1011, 0.1543,
                             0.1621, 0.1011, 0.11084, 0.104, 0.11426, 0.1406, 0.1079, 0.1035],
                            [0.1245, 0.1348, 0.2324, 0.1465, 0.249, 0.1592, 0.1533, 0.2354,
                             0.2461, 0.1533, 0.168, 0.1582, 0.1738, 0.2129, 0.1631, 0.1572],
                            [0.0996, 0.1074, 0.1865, 0.1167, 0.1992, 0.127, 0.12256, 0.1875,
                             0.1973, 0.12305, 0.1338, 0.126, 0.1387, 0.1699, 0.1309, 0.126],
                            [0.1318, 0.1426, 0.2461, 0.1543, 0.2637, 0.168, 0.1621, 0.248,
                             0.2617, 0.1621, 0.1777, 0.167, 0.1836, 0.2246, 0.1729, 0.166],
                            [0.1338, 0.1445, 0.25, 0.1572, 0.2676, 0.1699, 0.1641, 0.252,
                             0.2637, 0.165, 0.1807, 0.169, 0.1865, 0.2285, 0.1748, 0.169],
                            [0.1299, 0.1406, 0.2432, 0.1523, 0.2598, 0.166, 0.1602, 0.2451,
                             0.2578, 0.1602, 0.1758, 0.165, 0.1816, 0.2227, 0.1709, 0.1641],
                            [0.1196, 0.1299, 0.2236, 0.1406, 0.2402, 0.1533, 0.1475, 0.2256,
                             0.2373, 0.1475, 0.1611, 0.1514, 0.167, 0.2051, 0.1572, 0.1514]], dtype=np.float16),
                  np.array([[0.1172, 0.1602, 0.1943, 0.1113, 0.1719, 0.1865, 0.1709, 0.1514,
                             0.1484, 0.1504, 0.1533, 0.1309, 0.207, 0.1367, 0.1953, 0.1904],
                            [0.1172, 0.1602, 0.1943, 0.1113, 0.1719, 0.1865, 0.1709, 0.1514,
                             0.1484, 0.1504, 0.1533, 0.1309, 0.207, 0.1367, 0.1953, 0.1904],
                            [0.1172, 0.1602, 0.1943, 0.1113, 0.1719, 0.1865, 0.1709, 0.1514,
                             0.1484, 0.1504, 0.1533, 0.1309, 0.207, 0.1367, 0.1953, 0.1904],
                            [0.1172, 0.1602, 0.1943, 0.1113, 0.1719, 0.1865, 0.1709, 0.1514,
                             0.1484, 0.1504, 0.1533, 0.1309, 0.207, 0.1367, 0.1953, 0.1904],
                            [0.1172, 0.1602, 0.1943, 0.1113, 0.1719, 0.1865, 0.1709, 0.1514,
                             0.1484, 0.1504, 0.1533, 0.1309, 0.207, 0.1367, 0.1953, 0.1904],
                            [0.1172, 0.1602, 0.1943, 0.1113, 0.1719, 0.1865, 0.1709, 0.1514,
                             0.1484, 0.1504, 0.1533, 0.1309, 0.207, 0.1367, 0.1953, 0.1904],
                            [0.1172, 0.1602, 0.1943, 0.1113, 0.1719, 0.1865, 0.1709, 0.1514,
                             0.1484, 0.1504, 0.1533, 0.1309, 0.207, 0.1367, 0.1953, 0.1904],
                            [0.1172, 0.1602, 0.1943, 0.1113, 0.1719, 0.1865, 0.1709, 0.1514,
                             0.1484, 0.1504, 0.1533, 0.1309, 0.207, 0.1367, 0.1953, 0.1904],
                            [0.1172, 0.1602, 0.1943, 0.1113, 0.1719, 0.1865, 0.1709, 0.1514,
                             0.1484, 0.1504, 0.1533, 0.1309, 0.207, 0.1367, 0.1953, 0.1904],
                            [0.1172, 0.1602, 0.1943, 0.1113, 0.1719, 0.1865, 0.1709, 0.1514,
                             0.1484, 0.1504, 0.1533, 0.1309, 0.207, 0.1367, 0.1953, 0.1904],
                            [0.1172, 0.1602, 0.1943, 0.1113, 0.1719, 0.1865, 0.1709, 0.1514,
                             0.1484, 0.1504, 0.1533, 0.1309, 0.207, 0.1367, 0.1953, 0.1904],
                            [0.1172, 0.1602, 0.1943, 0.1113, 0.1719, 0.1865, 0.1709, 0.1514,
                             0.1484, 0.1504, 0.1533, 0.1309, 0.207, 0.1367, 0.1953, 0.1904],
                            [0.1172, 0.1602, 0.1943, 0.1113, 0.1719, 0.1865, 0.1709, 0.1514,
                             0.1484, 0.1504, 0.1533, 0.1309, 0.207, 0.1367, 0.1953, 0.1904],
                            [0.1172, 0.1602, 0.1943, 0.1113, 0.1719, 0.1865, 0.1709, 0.1514,
                             0.1484, 0.1504, 0.1533, 0.1309, 0.207, 0.1367, 0.1953, 0.1904],
                            [0.1172, 0.1602, 0.1943, 0.1113, 0.1719, 0.1865, 0.1709, 0.1514,
                             0.1484, 0.1504, 0.1533, 0.1309, 0.207, 0.1367, 0.1953, 0.1904],
                            [0.1172, 0.1602, 0.1943, 0.1113, 0.1719, 0.1865, 0.1709, 0.1514,
                             0.1484, 0.1504, 0.1533, 0.1309, 0.207, 0.1367, 0.1953, 0.1904]], dtype=np.float16)]
    else:
        raise ValueError(f"Unsupported device_type: '{device}'. Expected 'cpu' or 'gpu'.")

    return output
